"""
JumpServer轻量级SSH审计模块
基于JumpServer的核心审计功能，提供SSH会话录制和命令审计
"""

import os
import json
import time
import uuid
import threading
from datetime import datetime
from typing import Dict, List, Optional, Any
from django.conf import settings
from django.utils import timezone
from django.contrib.auth.models import User
from .models import Host, SSHSession, SSHCommand, SSHReplayData
import logging

logger = logging.getLogger(__name__)


class JumpServerAuditService:
    """JumpServer风格的审计服务"""

    def __init__(self):
        self.audit_dir = getattr(settings, 'JUMPSERVER_AUDIT_DIR', '/tmp/jumpserver_audit')
        self.ensure_audit_dir()

    def ensure_audit_dir(self):
        """确保审计目录存在"""
        if not os.path.exists(self.audit_dir):
            os.makedirs(self.audit_dir, mode=0o755)

    def create_session(self, user: User, host: Host, client_ip: str = None,
                      user_agent: str = '') -> SSHSession:
        """创建SSH会话记录"""
        session_id = str(uuid.uuid4())

        # 创建会话记录
        session = SSHSession.objects.create(
            session_id=session_id,
            user=user,
            host=host,
            client_ip=client_ip or '127.0.0.1',
            user_agent=user_agent,
            status='active'
        )

        # 创建审计文件目录
        date_str = datetime.now().strftime('%Y%m%d')
        session_dir = os.path.join(self.audit_dir, date_str, session_id)
        os.makedirs(session_dir, exist_ok=True)

        # 设置文件路径
        session.log_file_path = os.path.join(session_dir, 'session.log')
        session.replay_file_path = os.path.join(session_dir, 'replay.cast')
        session.save()

        logger.info(f"创建SSH会话: {session_id} - {user.username}@{host.hostname}")
        return session

    def close_session(self, session_id: str, exit_code: int = 0):
        """关闭SSH会话"""
        try:
            session = SSHSession.objects.get(session_id=session_id)
            session.end_time = timezone.now()
            session.status = 'closed' if exit_code == 0 else 'error'

            # 计算持续时间
            if session.start_time:
                duration = (session.end_time - session.start_time).total_seconds()
                session.duration = int(duration)

            # 统计命令数量
            session.command_count = session.commands.count()
            session.save()

            logger.info(f"关闭SSH会话: {session_id}")

        except SSHSession.DoesNotExist:
            logger.error(f"会话不存在: {session_id}")

    def log_command(self, session_id: str, command: str, command_type: str = 'input',
                   output: str = '', working_directory: str = '',
                   duration: float = 0, exit_code: int = None) -> Optional[SSHCommand]:
        """记录SSH命令"""
        try:
            session = SSHSession.objects.get(session_id=session_id)

            # 风险评估
            risk_level, is_dangerous = self._assess_command_risk(command)

            # 创建命令记录
            cmd_record = SSHCommand.objects.create(
                session=session,
                command_type=command_type,
                command=command,
                output=output,
                duration=duration,
                exit_code=exit_code,
                risk_level=risk_level,
                is_dangerous=is_dangerous,
                working_directory=working_directory
            )

            # 更新会话命令计数
            session.command_count = session.commands.count()
            session.save()

            # 写入日志文件
            self._write_command_log(session, cmd_record)

            return cmd_record

        except SSHSession.DoesNotExist:
            logger.error(f"会话不存在: {session_id}")
            return None

    def log_replay_data(self, session_id: str, data_type: str, data: str,
                       timestamp: float = None, terminal_width: int = 80,
                       terminal_height: int = 24) -> Optional[SSHReplayData]:
        """记录回放数据"""
        try:
            session = SSHSession.objects.get(session_id=session_id)

            if timestamp is None:
                timestamp = time.time()

            # 计算相对时间戳
            if session.start_time:
                relative_timestamp = timestamp - session.start_time.timestamp()
            else:
                relative_timestamp = 0

            # 创建回放数据记录
            replay_data = SSHReplayData.objects.create(
                session=session,
                timestamp=relative_timestamp,
                data_type=data_type,
                data=data,
                terminal_width=terminal_width,
                terminal_height=terminal_height
            )

            # 写入回放文件
            self._write_replay_data(session, replay_data)

            return replay_data

        except SSHSession.DoesNotExist:
            logger.error(f"会话不存在: {session_id}")
            return None

    def _assess_command_risk(self, command: str) -> tuple:
        """评估命令风险等级"""
        command = command.strip().lower()

        # 严重风险命令
        critical_commands = [
            'rm -rf /', 'rm -rf /*', 'dd if=/dev/zero', 'mkfs',
            'fdisk', 'parted', 'format', ':(){ :|:& };:'
        ]

        # 高风险命令
        high_risk_commands = [
            'rm -rf', 'rm -f', 'chmod 777', 'chown -R',
            'iptables -F', 'systemctl stop', 'service stop',
            'kill -9', 'pkill', 'killall', 'reboot', 'shutdown',
            'passwd', 'userdel', 'groupdel'
        ]

        # 中风险命令
        medium_risk_commands = [
            'rm ', 'mv ', 'cp -f', 'chmod', 'chown',
            'systemctl', 'service', 'crontab', 'at ',
            'sudo', 'su ', 'wget', 'curl', 'nc ', 'netcat'
        ]

        # 检查严重风险
        for cmd in critical_commands:
            if cmd in command:
                return 'critical', True

        # 检查高风险
        for cmd in high_risk_commands:
            if command.startswith(cmd) or f' {cmd}' in command:
                return 'high', True

        # 检查中风险
        for cmd in medium_risk_commands:
            if command.startswith(cmd) or f' {cmd}' in command:
                return 'medium', False

        return 'low', False

    def _write_command_log(self, session: SSHSession, command: SSHCommand):
        """写入命令日志文件"""
        if not session.log_file_path:
            return

        try:
            log_entry = {
                'timestamp': command.timestamp.isoformat(),
                'command_type': command.command_type,
                'command': command.command,
                'output': command.output,
                'duration': command.duration,
                'exit_code': command.exit_code,
                'risk_level': command.risk_level,
                'is_dangerous': command.is_dangerous,
                'working_directory': command.working_directory
            }

            with open(session.log_file_path, 'a', encoding='utf-8') as f:
                f.write(json.dumps(log_entry, ensure_ascii=False) + '\n')

        except Exception as e:
            logger.error(f"写入命令日志失败: {e}")

    def _write_replay_data(self, session: SSHSession, replay_data: SSHReplayData):
        """写入回放数据文件 (asciinema格式)"""
        if not session.replay_file_path:
            return

        try:
            # asciinema格式: [timestamp, event_type, data]
            event_type = 'i' if replay_data.data_type == 'input' else 'o'
            cast_entry = [replay_data.timestamp, event_type, replay_data.data]

            with open(session.replay_file_path, 'a', encoding='utf-8') as f:
                f.write(json.dumps(cast_entry, ensure_ascii=False) + '\n')

        except Exception as e:
            logger.error(f"写入回放数据失败: {e}")

    def get_session_replay_data(self, session_id: str) -> Dict[str, Any]:
        """获取会话回放数据"""
        try:
            logger.info(f"🔍 获取回放数据，会话ID: {session_id}")
            session = SSHSession.objects.get(session_id=session_id)
            logger.info(f"🔍 找到会话: {session.user.username}@{session.host.hostname}")

            # 从数据库获取回放数据
            replay_data = session.replay_data.all().order_by('timestamp')
            logger.info(f"🔍 回放数据条数: {replay_data.count()}")

            replay_list = [{
                'timestamp': data.timestamp,
                'data_type': data.data_type,
                'data': data.data
            } for data in replay_data]

            # 获取命令列表
            commands = session.commands.filter(command_type='input').order_by('timestamp')
            logger.info(f"🔍 输入命令条数: {commands.count()}")

            command_list = []
            for cmd in commands:
                try:
                    command_list.append({
                        'timestamp': cmd.timestamp.isoformat(),
                        'command': cmd.command,
                        'risk_level': cmd.risk_level,
                        'risk_level_display': cmd.get_risk_level_display(),
                        'duration': cmd.duration,
                        'is_dangerous': cmd.is_dangerous
                    })
                except Exception as e:
                    logger.error(f"❌ 处理命令记录失败: {e}, 命令: {cmd.command}")

            # 会话信息
            session_info = {
                'session_id': session.session_id,
                'user': session.user.username,
                'host': session.host.hostname,
                'host_ip': session.host.public_ip,
                'client_ip': session.client_ip,
                'start_time': session.start_time.isoformat() if session.start_time else None,
                'end_time': session.end_time.isoformat() if session.end_time else None,
                'duration': session.duration or 0,
                'status': session.status,
                'command_count': session.command_count or 0
            }

            # 计算总时长
            total_time = replay_data.last().timestamp if replay_data.exists() else session.duration or 0

            logger.info(f"✅ 回放数据获取成功，回放条数: {len(replay_list)}, 命令条数: {len(command_list)}")

            return {
                'session_info': session_info,
                'replay_data': replay_list,
                'commands': command_list,
                'total_time': total_time
            }

        except SSHSession.DoesNotExist:
            logger.error(f"❌ 会话不存在: {session_id}")
            return None
        except Exception as e:
            logger.error(f"❌ 获取回放数据失败: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")
            return None


# 全局JumpServer审计服务实例
jumpserver_audit = JumpServerAuditService()
